import re
import datetime

import pandas as pd
from django.db import models, connection


def date_format(raw_date, fmt):
    '''
    函数：日期格式化
        raw_date: 原始日期（支持 datetime.datetime, datetime.date, str 类型数据）
        fmt: 输出日期类型（支持 datetime.datetime, datetime.date, str）
    '''
    assert isinstance(raw_date, (datetime.datetime, datetime.date, str)), \
        "Error: Type %s is not accepted."%type(raw_date)
    if isinstance(raw_date, (datetime.date, str)):
        raw_date = str(raw_date)
        if re.findall(r'\d+-\d+-\d+', raw_date):
            raw_date = datetime.datetime.strptime(
                re.findall(r'\d+-\d+-\d+', raw_date)[0],
                "%Y-%m-%d"
            )
        elif len(raw_date) == 8:
            raw_date = datetime.datetime.strptime(raw_date, "%Y%m%d")
        else:
            assert raw_date, "Error: Date format is not accepted."
    
    if fmt == datetime.datetime:
        return raw_date
    elif fmt == datetime.date:
        return raw_date.date()
    elif fmt == str:
        return str(raw_date.date())


def date_diff(start_date, end_date):
    '''
    函数：日期差
        start_date: 开始日期（2022-01-01）
        end_date: 结束日期（2023-01-01）
    '''
    return abs((date_format(end_date, datetime.date) - date_format(start_date, datetime.date)).days)


def date_offset(raw_date, offset, fmt=datetime.date):
    '''
    函数：日期偏移
        raw_date: 原始日期（支持 datetime.datetime, datetime.date, str 类型数据）
        offset: 偏移量 （支持 整形 int）
        fmt: 输出日期类型（支持 datetime.datetime, datetime.date, str）
    '''
    date = date_format(raw_date, datetime.date)
    date += datetime.timedelta(days=offset)
    return date_format(date, fmt)


def is_weekend(raw_date):
    '''
    函数：日期是否位于周末
        raw_date: 原始日期（支持 datetime.datetime, datetime.date, str 类型数据）
    返回：
        boolean
    '''
    date = date_format(raw_date, datetime.date)
    return date.weekday() == 4


def get_last_month(date):
    '''
    函数：获取上月月份
        date: 原始日期（支持 datetime.datetime, datetime.date, str 类型数据）
    返回：
        str: '2022-02'
    '''
    raw_date = date_format(date, str)
    this_month = datetime.datetime.strptime(raw_date[:-3], '%Y-%m')
    last_month = this_month - datetime.timedelta(days = 1)
    return last_month.strftime('%Y-%m')


def get_pre_quarter(date, order=1):
    '''
    函数：获取上季度月份
        date:   原始日期（支持 datetime.datetime, datetime.date, str 类型数据）
        order:  前第 N 个季度
    返回：
        str: '2022-02'
    '''
    raw_date = date_format(date, datetime.date)
    quarters = [
        f'{year}-{month:02d}'
        for year in (raw_date.year, raw_date.year - 1)
        for month in (3, 6, 9, 12)
    ]
    quarters = sorted([q for q in quarters if q < raw_date.strftime('%Y-%m')])
    
    if order <= len(quarters):
        return quarters[-order]

    
def get_quarter_periods():
    '''
    函数：获取季度周期清单
    返回：
        [
            datetime.date(2022, 3, 31),
            datatime.date(2021, 12, 31),
            ...
            datetime.date(2022, 3, 31)
        ]
    '''
    return [
        datetime.date(year, month, date).strftime('%Y%m%d').replace('0331', 'Q1').replace('0630', 'Q2').replace('0930', 'Q3').replace('1231', 'Y')
        for year in range(datetime.date.today().year, 1999, -1)
        for month, date in [(12, 31), (9, 30), (6, 30), (3, 31)]
        if datetime.date(year, month, date) < datetime.date.today()
    ]


def tushare_batch_fetch(function, *args, **kwargs):
    '''
    函数：Tushare 批量拉取
    参数：
        function:   tushare 函数
        * args:     tushare 位置参数
        ** kwargs:  tushare 字典参数
    返回：
        dataframe
    '''
    batch_size = kwargs.pop('batch_size', 10 ** 5)
    page_count = kwargs.pop('page_count', 2 ** 10)

    res = None

    for epoch in range(page_count):
        kwargs.update({'limit': batch_size, 'offset': batch_size * epoch})
        df = function(*args, **kwargs)

        if res is None:
            res = df
        elif not df.empty:
            res = pd.concat([res, df])
        else:
            break

    return res
        


class TruncateManager(models.Manager):
    def truncate(self):
        with connection.cursor() as cursor:
            cursor.execute(f"TRUNCATE TABLE {self.model._meta.db_table}")